/*
* Copyright (C) 2020 ~ 2024 chaigec All Rights Reserved.
*
* Author:     chaigec <chaigec@163.com>
* Maintainer: chaigec <chaigec@163.com>

 * Use of this source code is governed by MIT license that can be found in the
 * LICENSE file in the root of the source tree. All contributing project authors
 * may be found in the AUTHORS file in the root of the source tree.
*/

#include "MkWebSocketMessage.h"
#include "MkUtil/MkUtil.h"
#include "MkUtil/MkLog.h"

CMkWebSocketMessage::CMkWebSocketMessage()
    : m_Fin(1)
    , m_Reserved(0)
    , m_Mask(0)
    , m_Opcode(TextFrame)
{

}

void CMkWebSocketMessage::SetMark(BOOL bMask)
{
    m_Mask = bMask;
    if (bMask) {
        MkString Tmp;
        CMkUtil::CreateRandomString(sizeof(Uint32), Tmp);
        memcpy(&m_MaskingKey, Tmp.c_str(), sizeof(Uint32));
    }
}

Uint32 CMkWebSocketMessage::AddContent(const Uint8* pBuf, Uint32 nLen, OUT Uint32& nUsed)
{
    if (nLen < 2) {
        nUsed = 0;
        return NotHasEnoughBuffer;
    }
    Uint8 *pOffset = (Uint8*)pBuf;
    Uint8 *pFirst = pOffset;
    m_Fin = *pOffset >> 7;
    m_Reserved = (*pOffset >> 4) & 0x07;
    m_Opcode = static_cast<MkWebSocketFrameType>((*pOffset) & 0x0f);
    pOffset++;
    m_Mask = *pOffset >> 7;
    m_PayloadLenth = *pOffset & 0x7F;
    //2 headers
    if (nLen < m_PayloadLenth + 2) {
        nUsed = 0;
        return NotHasEnoughBuffer;
    }
    pOffset++;
    if (126 == m_PayloadLenth) {
        if (nLen < 4) {
            nUsed = 0;
            return NotHasEnoughBuffer;
        }
        m_PayloadLenth = *pOffset << 8;
        pOffset++;
        m_PayloadLenth += *pOffset;
        pOffset++;
    } else if (127 == m_PayloadLenth) {
        if (nLen < 10) {
            nUsed = 0;
            return NotHasEnoughBuffer;
        }
        m_PayloadLenth = (Uint64)*pOffset << 56;
        pOffset++;
        m_PayloadLenth += (Uint64)*pOffset << 48;
        pOffset++;
        m_PayloadLenth += (Uint64)*pOffset << 40;
        pOffset++;
        m_PayloadLenth += (Uint64)*pOffset << 32;
        pOffset++;
        m_PayloadLenth += (Uint64)*pOffset << 24;
        pOffset++;
        m_PayloadLenth += (Uint64)*pOffset << 16;
        pOffset++;
        m_PayloadLenth += *pOffset << 8;
        pOffset++;
        m_PayloadLenth += *pOffset;
        pOffset++;
    }

    if (1 == m_Mask) {
        if (pOffset - pFirst + m_PayloadLenth + 4 > nLen) {
            nUsed = 0;
            return NotHasEnoughBuffer;
        }
        m_MaskingKey = *pOffset;
        pOffset++;
        m_MaskingKey += *pOffset << 8;
        pOffset++;
        m_MaskingKey += *pOffset << 16;
        pOffset++;
        m_MaskingKey += *pOffset << 24;
        pOffset++;

        for (Uint64 i = 0; i < m_PayloadLenth; i++) {
            pOffset[i] = pOffset[i] ^ ((unsigned char *)(&m_MaskingKey))[i % 4];
        }
    }
    nUsed = static_cast<Uint32>(pOffset - pFirst + m_PayloadLenth);
    if (nUsed > nLen) {
        return NotHasEnoughBuffer;
    }
    m_pPayloadData = (Int8*)pOffset;
    return NoneError;
}

Uint32 CMkWebSocketMessage::PacketHeader(Uint64 PayloadLen, MkString& Header)
{
    Uint8 pHeader[Len16] = { 0 };
    Uint8 *pOffset = pHeader;
    *pOffset = (m_Fin << 7) | (m_Reserved << 4) | m_Opcode;
    pOffset++;
    if (PayloadLen > 0xFFFF) {
        *pOffset = (m_Mask << 7) | 127;
        pOffset++;
        //next 8byte len
        *pOffset = (PayloadLen >> 56) & 0xff;
        pOffset++;
        *pOffset = (PayloadLen >> 48) & 0xff;
        pOffset++;
        *pOffset = (PayloadLen >> 40) & 0xff;
        pOffset++;
        *pOffset = (PayloadLen >> 32) & 0xff;
        pOffset++;
        *pOffset = (PayloadLen >> 24) & 0xff;
        pOffset++;
        *pOffset = (PayloadLen >> 16) & 0xff;
        pOffset++;
        *pOffset = (PayloadLen >> 8) & 0xff;
        pOffset++;
        *pOffset = PayloadLen & 0xFF;
        pOffset++;
    } else if (PayloadLen > 125) {
        *pOffset = (m_Mask << 7) | 126;
        pOffset++;
        //next 2byte
        *pOffset = (PayloadLen >> 8) & 0xff;
        pOffset++;
        *pOffset = PayloadLen & 0xFF;
        pOffset++;
    } else {
        *pOffset = (m_Mask << 7) | (PayloadLen & 0xFF);
        pOffset++;
    }
    if (m_Mask) {
        memcpy(pOffset, &m_MaskingKey, sizeof(Uint32));
        pOffset += sizeof(Uint32);

    }
    Header = MkString((Int8*)pHeader, pOffset - pHeader);
    return NoneError;
}

Uint32 CMkWebSocketMessage::PacketContent(const MkString& Content, MkString& Out)
{
    if (m_Mask) {
        for (size_t i = 0; i < Content.size(); i++) {
            Out += (Int8)(((Uint8)Content[i]) ^ ((unsigned char *)(&m_MaskingKey))[i % 4]);
        }
    } else {
        Out = Content;
    }
    return NoneError;
}

CMkWebSocketMessage::~CMkWebSocketMessage()
{

}